package dev.alm.chatbot.controller;

import dev.alm.chatbot.common.AiRequest;
import dev.alm.chatbot.common.AiResponse;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.HashMap;

@RestController
@RequestMapping("ai")
public class AiController {

    private final ChatClient chatClient;

    public AiController(ChatClient chatClient) {
        this.chatClient = chatClient;
    }

    @PostMapping("chat")
    public AiResponse chat(@RequestBody AiRequest aiRequest) {
        String content = chatClient.prompt().user(aiRequest.getPrompt()).call().content();
        // 把content中以<think>开头，一直到</think>结束的文本取出来
        String think = content.substring(content.indexOf("<think>") + 7, content.indexOf("</think>"));
        // 把</think>后面的内容取出来
        String answer = content.substring(content.indexOf("</think>") + 8);
        AiResponse aiResponse = new AiResponse();
        aiResponse.setCode(200);
        aiResponse.setMessage("success");
        HashMap<String, String> map = new HashMap<>();
        map.put("think", think);
        map.put("answer", answer);
        aiResponse.setData(map);
        return aiResponse;
    }

    @RequestMapping(value = "stream", produces = "text/html;charset=utf-8")
    public Flux<String> stream(String prompt) {
        return chatClient.prompt().user(prompt).stream().content();
    }

}
